package softsdk

import (
	"crypto/rand"
	"fmt"
	"math/big"
	"os"
	"strings"
)

// 安全参数
const (
	SystemLength = (1 << 16)
	AlphaLength  = 80
	RandLength   = 80
)

const FheByteLen = 128 // 一个同态密文有128字节

var (
	pFheX0    *big.Int
	pFheEncX1 *big.Int
	pFheEncX2 *big.Int
	pFheSK    *big.Int
)

func FheSoftEncUtf8(utf8Text string) ([]byte, error) {
	// 本函数输出的密文是大端存储，因为FillBytes方法得到的字节串是大端存储。
	uint16Text := []uint16{}
	for i := 0; i < len(utf8Text)/2; i++ {
		b1 := utf8Text[i*2]
		b2 := utf8Text[i*2+1]
		uint16Elem := (uint16(b1) << 8) | uint16(b2)
		uint16Text = append(uint16Text, uint16Elem)
	}
	if len(utf8Text)%2 != 0 {
		b1 := utf8Text[len(utf8Text)-1]
		b2 := 0
		uint16Elem := (uint16(b1) << 8) | uint16(b2)
		uint16Text = append(uint16Text, uint16Elem)
	}
	utf8EncText := []byte{}
	for _, uint16Elem := range uint16Text {
		uint16Enc, err := fheSoftEncCore(uint16Elem)
		if err != nil {
			return nil, err
		}
		// padding enc to FheByteLen bit
		uint16EncBytes := make([]byte, FheByteLen)
		uint16Enc.FillBytes(uint16EncBytes) // uint16EncBytes是大端存储
		utf8EncText = append(utf8EncText, uint16EncBytes...)
	}
	return utf8EncText, nil
}

func FheSoftDecUtf8(utf8EncText []byte) (string, error) {
	if len(utf8EncText)%FheByteLen != 0 {
		return "", fmt.Errorf("EncText length is not a multiple of %d", FheByteLen)
	}
	utf8DecTextBytes := []byte{}
	for i := 0; i < len(utf8EncText)/FheByteLen; i++ {
		uint16EncText := utf8EncText[i*FheByteLen : (i+1)*FheByteLen]
		uint16ClearText, err := fheSoftDecCore(new(big.Int).SetBytes(uint16EncText))
		if err != nil {
			return "", err
		}
		b2 := byte(uint16ClearText & 0xFF)
		b1 := byte(uint16ClearText >> 8)
		utf8DecTextBytes = append(utf8DecTextBytes, b1, b2)
	}
	var str = string(utf8DecTextBytes)
	if len(str) > 0 && int(str[len(str)-1]) == 0 {
		str = str[:len(str)-1]
	}
	return str, nil
}

// FheSoftSub 两个同态密文相减，得到同态密差
func FheSoftSub(enc1, enc2 []byte) (diff []byte, err error) {
	n1 := len(enc1)
	n2 := len(enc2)
	if n1 != n2 {
		return nil, fmt.Errorf("enc1 and enc2 lengths not equal")
	}
	if n1 == 0 {
		return nil, fmt.Errorf("enc1 and enc2 length is 0")
	}
	if n1%FheByteLen != 0 {
		return nil, fmt.Errorf("enc1 and enc2 lengths is not a multiple of %d", FheByteLen)
	}
	n := n1
	for i := 0; i < n/FheByteLen; i++ {
		num1 := new(big.Int).SetBytes(enc1[i*FheByteLen : (i+1)*FheByteLen])
		num2 := new(big.Int).SetBytes(enc2[i*FheByteLen : (i+1)*FheByteLen])
		if num1.Cmp(num2) == -1 { // num1 < num2
			num1.Add(num1, pFheX0)
		}
		subNum := bigIntSub(num1, num2)
		subBytes := make([]byte, FheByteLen)
		subNum.FillBytes(subBytes)
		diff = append(diff, subBytes...)
	}
	return diff, nil
}

// FheSoftDecDiff 解密同态密差，若解密结果为0返回true
func FheSoftDecDiff(diff []byte) (bool, error) {
	n := len(diff)
	if n == 0 {
		return false, fmt.Errorf("diff length is 0")
	}
	if n%FheByteLen != 0 {
		return false, fmt.Errorf("diff length is not a multiple of %d", FheByteLen)
	}
	for i := 0; i < n/FheByteLen; i++ {
		dec, err := fheSoftDecCore(new(big.Int).SetBytes(diff[i*FheByteLen : (i+1)*FheByteLen]))
		if err != nil {
			return false, err
		}
		if dec != 0 {
			return false, nil
		}
	}
	return true, nil
}

func readSoftEncPK(filePath string) error {
	pkBytes, err := os.ReadFile(filePath)
	if err != nil {
		return fmt.Errorf("unable to open fhe public key file")
	}
	pkStr := strings.Replace(string(pkBytes), "\r", "", -1)
	pkLines := strings.Split(pkStr, "\n")
	var pks []*big.Int
	for _, line := range pkLines {
		if line == "" {
			continue
		}
		pkXi, errBool := new(big.Int).SetString(line, 0)
		if !errBool {
			panic(line + " is not hexadecimal")
		}
		pks = append(pks, pkXi)
	}
	pFheX0 = pks[0]
	pFheEncX1 = pks[1] // random pk[i]
	pFheEncX2 = pks[6] // random pk[i]
	if pFheX0.Cmp(big.NewInt(0)) == 0 || pFheEncX1.Cmp(big.NewInt(0)) == 0 || pFheEncX2.Cmp(big.NewInt(0)) == 0 {
		return fmt.Errorf("error: fhe_x0、fhe_enc_x1 and fhe_enc_x2 loading failed")
	}
	return nil
}

func readSoftEncSK(filePath string) error {
	skBytes, err := os.ReadFile(filePath)
	if err != nil {
		return fmt.Errorf("unable to open fhe secret key file")
	}
	skStr := strings.Replace(string(skBytes), "\r", "", -1)
	skLines := strings.Split(skStr, "\n")
	skX0, errBool := new(big.Int).SetString(skLines[0], 0)
	if !errBool {
		panic(skLines[0] + " is not hexadecimal")
	}
	pFheSK = skX0
	if pFheSK.Cmp(big.NewInt(0)) == 0 {
		return fmt.Errorf("error: fhe_sk loading failed")
	}
	return nil
}

func fheSoftEncCore(plainText uint16) (*big.Int, error) {
	if pFheX0 == nil || pFheEncX1 == nil || pFheEncX2 == nil {
		return nil, fmt.Errorf("error: fhe_x0 or fhe_enc_x1 or fhe_enc_x2 not load")
	}
	plainTextBigInt := big.NewInt(int64(plainText))
	systemLengthBigInt := big.NewInt(int64(SystemLength))
	subsetProduct := bigIntMul(
		bigIntAdd(pFheEncX1, pFheEncX2),
		nfheRand(AlphaLength),
	)
	return bigIntMod(
		bigIntAdd(
			plainTextBigInt,
			bigIntMul(systemLengthBigInt, nfheRand(RandLength)),
			bigIntMul(systemLengthBigInt, subsetProduct),
		),
		pFheX0,
	), nil
}

func fheSoftDecCore(encText *big.Int) (uint16, error) {
	if pFheSK == nil {
		return 0, fmt.Errorf("error: fhe_sk not load")
	}
	preDec1 := bigIntMod(encText, pFheSK)
	preDec2 := bigIntAdd(
		bigIntSub(
			new(big.Int).Lsh(big.NewInt(1), 256),
			pFheSK),
		preDec1)
	cond := preDec1.Cmp(bigIntDiv(pFheSK, big.NewInt(2)))
	systemLengthBigInt := big.NewInt(SystemLength)
	if cond == -1 || cond == 0 { // preDec1 <= fheSk/2
		return uint16(bigIntMod(preDec1, systemLengthBigInt).Int64()), nil
	} else {
		return uint16(bigIntMod(preDec2, systemLengthBigInt).Int64()), nil
	}
}

func nfheRand(randLength int) *big.Int {
	randNum, err := rand.Int(rand.Reader, new(big.Int).Lsh(big.NewInt(1), uint(randLength)))
	if err != nil {
		panic("nfheRand rand.Int panic")
	}
	return randNum
}

func bigIntAdd(bigInts ...*big.Int) *big.Int {
	res := big.NewInt(0)
	for _, bigInt := range bigInts {
		res.Add(res, bigInt)
	}
	return res
}

func bigIntSub(bitInt1, bigInt2 *big.Int) *big.Int {
	return new(big.Int).Sub(bitInt1, bigInt2)
}

func bigIntMul(bitInt1, bigInt2 *big.Int) *big.Int {
	return new(big.Int).Mul(bitInt1, bigInt2)
}

func bigIntDiv(bitInt1, bigInt2 *big.Int) *big.Int {
	return new(big.Int).Div(bitInt1, bigInt2)
}

func bigIntMod(bigInt1 *big.Int, bigInt2 *big.Int) *big.Int {
	return new(big.Int).Mod(bigInt1, bigInt2)
}
